Add 8-wave warp-pipeline GEMM kernels (a16w16 / a8w8 / a4w4)#44
Open
zhanglx13 wants to merge 40 commits into
Open
Add 8-wave warp-pipeline GEMM kernels (a16w16 / a8w8 / a4w4)#44zhanglx13 wants to merge 40 commits into
zhanglx13 wants to merge 40 commits into
Conversation
Ports the 8-wave warp-pipeline GEMM from AMD-Triton/gluon-kernels into the
tutorial layout so the existing bench / rocprof / ATT tooling can drive it.
Contrast with the 4-wave a16w16 kernels: 8 warps (2x4), 256x256x32 tiles,
3-buffer (triple) LDS ring, warp_pipeline_stage wave-level scheduling. The
llir+amdgcnas toolchain is 4-wave-only (fails RA here); this kernel schedules
itself via warp_pipeline_stage.
Files:
- matmul_kernel.py 2x-unrolled kernel (gemm_async_warp_pipeline) + matmul wrappers
- matmul_kernel_unroll3.py 3x-unrolled variant: constant LDS buffer indices (no tile%3
address math), v9 XCD-aware PID remap (GROUP_SIZE_M=4)
- common.py PID mapping (get_pids / get_pid_m_n), store epilogue, async loads
- bench.py correctness + do_bench + rocprof rotating-tensor mode (--unroll {2,3})
- collect_perf.py rocprof kernel-trace TFLOPS + ATT MFMA efficiency
- collect_counters.py VMEM-latency / L2-miss hardware counters
Measured on MI350X at 4096x4096x8192 fp16 (no-AGPR via TRITON_HIP_AGPR_ALLOC=0,0):
3x + v9 XCD remap reaches ~909 TFLOPS rocprof, 85% MFMA efficiency, 0 spills.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: zhanglx13 <lixun.zhang@amd.com>
Each unrolled step placed cdna4_async.wait_group(2) AFTER the mfma region and issued the async global load (GR, buffer_load_to_shared) before the shared load (LR, ds_read) in the mem region. With wait_group after the mfma, the per-iteration s_barrier synchronizes only the 8 waves' EXECUTION -- it does not wait for the async GR issued by the other 4-wave group to land. A wave could therefore consume an LDS tile whose filling copy had not completed (a cross-wave RAW on the 3-buffer LDS ring). The kernel passed the allclose check by timing luck, not by construction. Move wait_group(2) before the mfma region and issue LR before GR in all three steps. This drains the GR vmcnt ahead of the mfma, gating GR completion before the data is consumed -> race-free by construction. (LR-first alone, with wait_group still after the mfma, is miscompiled to a malformed bare s_barrier between LR and GR; moving wait_group ahead of the mfma is what lets the framework lower the reordered region correctly.) Loads stay relaxed (load_shared_relaxed): the no-alias hint elides a redundant LDS barrier that the non-relaxed smem.load form inserts here (~6% slower). Perf (4096x4096x8192 fp16, MI350X, no-AGPR, rocprof cold-rotating): ~887-899 TFLOPS / 77.8% per-SIMD MFMA (cold) / 202 VGPR / 0 spill, vs the prior racy layout's 909 / 85.2% / 198 -- ~2% throughput for guaranteed correctness. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: zhanglx13 <lixun.zhang@amd.com>
The prior commit (race fix) placed step-0's cdna4_async.wait_group(2) right before the mfma stage, but the loop body begins with the loop-index arithmetic (base_tile/pf*). WarpPipeliner.createPipeline() rejects pipelining when it meets an AsyncWaitOp (wait_group) while a stage cluster is non-empty -- the index arith populated the cluster first, so step-0's wait_group tripped the bailout and silently disabled the warp-pipeline: 0 s_setprio, 3 s_barriers/iter instead of 6, and the scheduler split the mfma region (1 mfma hoisted above the vmcnt drain). Hoist step-0's wait_group above the index arithmetic so it lands at an empty cluster boundary; pipelining is preserved. It still precedes step-0's mfma (race-free), and wait_group is a fence so moving it above pure arithmetic is semantically free. Steps 1-2 keep wait_group inline (each follows a mem-stage border = empty cluster, so no rejection there). Perf (4096x4096x8192 fp16, MI350X, no-AGPR, rocprof cold-rotating): 913.8 TFLOPS / 84.8% per-SIMD MFMA / 202 VGPR / 0 spill -- warp-pipeline back on (s_setprio + 2 s_barriers/step, 32-mfma region unsplit), vs the regressed 887-899 / 77.8% from the pipeline-disabled commit. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: zhanglx13 <lixun.zhang@amd.com>
…yle) Two refactors to match the a16w16/v9 kernel conventions: 1. Local loads use the non-relaxed smem.index(i).load(layout) form instead of cdna4_async.load_shared_relaxed(...). The non-relaxed form re-inserts a redundant `s_waitcnt lgkmcnt(0); s_barrier` between the ds_read (LR) and the buffer_load (GR) in each mem region (s_barrier 17->22). That drops warm MFMA eff 84.5%->76.8%, but only ~0.6% on the cold-rotating rocprof TFLOPS -- the real workload is memory-bound, so the extra compute barrier is largely hidden. 2. issue_async_load_a/b helpers are inlined as direct buffer_load_to_shared calls. Offsets are precomputed once (tile-relative K=0) and the absolute K index is carried in the base pointer (base + k*BLOCK_K*stride), the v9 style. Perf-neutral; reduces VGPR pressure 196->190 (0 spills). Net (4096x4096x8192 fp16, MI350X, no-AGPR, rocprof cold-rotating, upstream triton 3.8.0): 900.1 TFLOPS / 76.8% warm MFMA / 190 VGPR / 0 spill, vs the relaxed+helper baseline 905.7 / 84.5% / 196. Warp-pipeline stays on (s_setprio present); correctness verified. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: zhanglx13 <lixun.zhang@amd.com>
Replace the per-call base arithmetic (a_base + pf*BLOCK_K*stride at each buffer_load_to_shared) with the a16w16/v9 pointer-walk style: - Precompute 3 distinct offsets per operand (a_offs0/1/2, b_offs0/1/2) for the triple's K = 0,1,2 relative to a running base. - The 3 unrolled mem regions all prefetch off the SAME a_base/b_base with those offsets; the base is advanced by one triple (a_base += 3*BLOCK_K*stride_ak) after the prologue and in step 2 of each loop iteration. - Drops the loop-index arithmetic (base_tile/pf*), so wait_group(2) is now naturally the first op in the loop body -- no hoisting needed to keep WarpPipeliner from rejecting; the warp-pipeline stays on. - The tail reuses the walked base + a_offs0/1. Correct across all benchmarked K (512..16384, both tail branches exercised). Perf (4096x4096x8192 fp16, MI350X, no-AGPR, rocprof cold-rotating, upstream triton 3.8.0): 903.7 TFLOPS / 79.5% MFMA / 188 VGPR / 0 spill, vs the previous per-call-base layout's 900.1 / 76.8% / 190 -- +3.6 TFLOPS, +2.6pp MFMA, -2 VGPR. The constant offset deltas fold cheaply (no register blowup from 3 offset sets). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: zhanglx13 <lixun.zhang@amd.com>
…rnels
Two things in one cleanup pass:
1) Relaxed local_load. The mem region does smemA.index(k+1).load (LR) then
buffer_load_to_shared(smemA.index(k)) (GR) -- same allocation, different
index. Triton's membar analysis can't disambiguate MemDescIndexOp sub-buffers
(it only reads static offsets from MemDescSubsliceOp), so it inserts a
redundant lgkmcnt(0)+s_barrier between LR and GR. Switching to
cdna4_async.load_shared_relaxed lets the AMD membarFilter skip it via the
async-wait token: s_barrier 22->17, MFMA 79.5%->85.1%, 912.7 TFLOPS at the
same 188 VGPR (4096^2 K8192 fp16, MI350X, no-AGPR). (The 4-wave v9 kernel
dodges the same barrier by using separate per-quadrant allocations.)
2) Versioned layout, mirroring a16w16/v0_naive..v9_beyond_hotloop:
- remove the old 2x matmul_kernel.py;
- matmul_kernel_unroll3.py -> v0_BK32_nS3/matmul_kernel.py (the relaxed
pointer-walk kernel above); fn + KERNEL_NAME renamed to match the dir;
- add v1_sliceMN_BK64_nS2/matmul_kernel.py -- a runnable copy of v0 scaffolded
for the next variant (M/N quadrant slicing a la v8_sliceMN, BLOCK_K=64,
2-buffer nS2); see its module docstring for the TODO;
- common.py: keep get_pids + store_result, drop the now-unused get_pid_m_n
and issue_async_load_a/b helpers (kernels inline buffer_load_to_shared);
- bench.py / collect_perf.py / collect_counters.py: --unroll {2,3} ->
--version {0,1} (VERSION_MAP + importlib); README rewritten for the layout.
Both versions verified correct (bench.py --version 0/1).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: zhanglx13 <lixun.zhang@amd.com>
Replace the v1 scaffold (was a copy of v0) with the M/N-sliced kernel from
/data/8wave-warp-pipeline-sliceMN.png. The 256x256 C tile is split into 4x
[128x128] quadrants (C_tl/bl/tr/br = A_top/bot x B_left/right), each quadrant
its OWN double-buffered LDS allocation (smemA_top/bot, smemB_left/right). Because
the four allocations have distinct buffer IDs, Triton's membar disambiguates the
local_load (LR) vs buffer_load_to_shared (GR) with no extra barrier, so the loads
stay NON-relaxed (the v9 trick) - no load_shared_relaxed needed here.
K loop unrolled 2x -> 8 mfma regions + 8 mem regions, each wrapped in
warp_pipeline_stage("mfma"/"mem") with cdna4_async.wait_group(5) hoisted BEFORE
each mfma (per the design image). 8-wave global-load layouts = v9's 4-wave layouts
plus one extra warp dim tiling M (A) / N (B); shared/dot/mfma layouts reused from
v9. B is pre-transposed to (N,K) and fed as logical (K,N) via strides.
Perf @4096^2 x8192 fp16, MI350X, no-AGPR (rocprof, cold rotating):
~1004 TFLOPS, +10% over v0 (~912). In-loop MFMA ~99.8% per-SIMD; the loop body
carries zero spill traffic (all 67 spills land in prologue/epilogue). Correct
across K 512..16384 for both fp16 and bf16.
Run: TRITON_HIP_AGPR_ALLOC="0,0" python bench.py --version 1 --K 8192 --dtype fp16
Signed-off-by: zhanglx13 <lixun.zhang@amd.com>
…eave) The 67 spills were all epilogue-internal -- the four live f32 [128x128] accumulators getting evicted to scratch (the loop body itself has zero scratch traffic; in-loop MFMA stays ~99.8%). Two kernel-side changes drop them to 0: 1) Store-side pointer-walk. Replace the four separate [128x128] gStoreLayoutC offset tensors (~32 VGPR each) with ONE shared within-quadrant offset tensor + four SCALAR base pointers (c_tl/bl/tr/br_base = c_base + const). All four quadrants share the same internal structure, so only the base differs. 2) De-interleave the epilogue. The old code ran acc_tr/acc_br's final mfmas AFTER the first buffer_stores, pinning the ~96 VGPR of dot operands live through the store phase on top of the 128 VGPR of accumulators. Finish all four mfmas first so the operands die, leaving only the 4 accumulators (+ one in-flight convert) live during the stores. Result @4096^2 x8192 fp16, MI350X, no-AGPR (rocprof, cold rotating): 67 -> 0 spills, 256 -> 242 VGPR, epilogue 32240 -> 16660 cyc (-48%), ~1007 -> ~1039 TFLOPS (+3.5%). Loop unchanged (s_setprio pipeline still on, 99.84% loop MFMA). Correct across K 512..16384 for both fp16 and bf16. Signed-off-by: zhanglx13 <lixun.zhang@amd.com>
Match the a16w16 convention (one README per version subdir): - v0_BK32_nS3/README.md: the tuned baseline -- BLOCK_K=32 / 3-buffer ring / 3x-unroll / warp_pipeline_stage / no-AGPR; the relaxed-load membar trick, the pointer-walk addressing, and the v9 XCD remap. - v1_sliceMN_BK64_nS2/README.md: M/N quadrant slicing with four separate per-quadrant LDS allocs (non-relaxed loads, no membar barrier), BLOCK_K=64 / 2-buffer, the 8-region warp-pipeline, and the epilogue spill fix (store-side pointer-walk + de-interleave) that took it to 0 spills / ~1039 TFLOPS. Also refresh the top-level README: drop the stale "v1 is a scaffold" note, link both subdir READMEs, and fold v1 into the perf table (noting loop-only vs whole-kernel MFMA efficiency). Signed-off-by: zhanglx13 <lixun.zhang@amd.com>
Rewrite kernels/gemm/a16w16-8wave/README.md from scratch around four things: a v0/v1/v9 design-comparison table; v0 as a port of AMD-Triton/gluon-kernels' f16_gemm_warp_pipeline_gfx950.py plus the fixes that lifted its per-SIMD MFMA from ~36% to ~85% (no-AGPR, 3x unroll, relaxed-load membar dodge, XCD remap); v1 as v9's sliceMN hot loop combined with the 8-wave warp-pipeline; and a v0-vs-v1 perf table (TFLOPS + MFMA eff) at K=8192/16384/32768, collected no-AGPR cold-rotating. The table also shows v0 degrading at large K (the non-sliced ring's buffer-load stall) while v1 climbs. Add (4096, 4096, 32768) to bench.py's get_x_vals so --K 32768 is selectable (used for the new perf rows). Signed-off-by: zhanglx13 <lixun.zhang@amd.com>
Replace section 2's four bullet points with a five-row table (as-ported baseline + the four fixes) carrying Step / TFLOPS / MFMA (per-SIMD) / Optimization columns, so the ~36% per-wave starting point and the rocprof cold-rotating progression (760 -> 820 -> 839 -> 909 -> ~915) are visible at a glance. Signed-off-by: zhanglx13 <lixun.zhang@amd.com>
§4: new MI355X table — v0/v1 (8-wave, no-AGPR) + a16w16 v9 (gemm-4waves,force-agpr, amdgcnas), cold rotating-2048. v9 edges v1 on TFLOPS at all K (+2/+3/+1%); v1 holds ~99% loop MFMA-eff; v0 is buffer-load-stall bound at large K. §2: note that the fix-progression numbers (and the MI350X §4 table) are from MI350X.
v0 MFMA-eff (~83/65/60%) confirmed stable across 5 runs — not jitter. The gap vs the MI350X table (85/78/80%) is real: MI355X's faster MFMA outpaces the memory-bound v0's unchanged load latency, so TFLOPS rises while loop MFMA-eff drops. v1/v9 unaffected.
ATT timelines at K=8192: v1 and v9 reach near-solid MFMA utilization; v0 is HBM-latency bound (clustered buffer_loads stall the MFMA units).
…) for the v0 port
Replace the TRITON_HIP_AGPR_ALLOC env-var/compiler-hook mechanism with
Triton's built-in llvm_fn_attrs launch option. Both v0 and v1 now pass
llvm_fn_attrs=(("amdgpu-agpr-alloc", "0,0"),) at launch, so the no-AGPR
setting lives in the kernel and survives a Triton rebuild without patching
the backend compiler.
Verified at K=8192 (no env var set): v0 1185 TFLOPS / 83.0% MFMA,
v1 1457 TFLOPS / 99.8% MFMA, agpr_count=0 / 0 spills for both — matching
the prior env-var path.
Update README/docs (main, v0, v1) and collect_counters.py to drop all
TRITON_HIP_AGPR_ALLOC references.
Add docs/warp_pipelining.md, a reference for the theory behind the a16w16-8wave warp-pipeline kernels: the phase-shifted two-group schedule, why it raises MFMA utilization, the barrier/membar rules, and the stage-level LDS dependency model. Adapted (condensed and re-grounded in the tutorial kernels) from Jungwook Park's original write-up, with attribution. Wire it into discovery: a row in the root README doc index and a pointer from the a16w16-8wave README.
Port the a16w16-8wave/v1_sliceMN warp-pipeline solution (8 warps, 2 waves/SIMD, warp_pipeline_stage, no-AGPR) to 8-bit and 4-bit: - a8w8-8wave/v1_sliceMN_BK128_nS2: BF8, mfma_scaled e5m2, ~2841 TFLOPS @ K=8192, 99.7% loop MFMA eff; beats every 4-wave a8w8 config. - a4w4-8wave/v1_sliceMN_BK256_nS2: MXFP4 with e8m0 scale pipeline + load-side pointer-walk (spill-free loop); ~4142 TFLOPS @ K=32768, matches 4-wave base. Each dir ships bench.py, collect_perf.py, common.py, and READMEs.
…current build - gemm/README.md: add §5 "8-Wave Warp-Pipeline Variants" — design, kernel + perf comparison tables, run commands, and 8-wave vs 4-wave summary. - Refresh all 8-wave README perf/VGPR numbers to the current build (MI355X, Triton 3.8.0); consolidate a16w16-8wave §4 into one MI355X table. - a8w8: on the newer compiler the 4-wave llir+amdgcnas path improved (~2746 -> 3216) and now leads the 8-wave; update the "beats every 4-wave config" claim accordingly. - a4w4: the ROCm 7.2.4 ATT decoder now decodes the FP4 scaled-MFMA (~57%), so the MFMA-eff column is no longer N/A. - a16w16-8wave v0 VGPR 188 -> 196 (compiler-version difference). - ROADMAP: mark 8-wave warp-pipeline GEMM as done; fix a4w4-8wave K-label. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Lixun Zhang <lixun.zhang@amd.com>
…t build) Rework the Performance Summary into a 4-wave vs 8-wave table for all four data types on the current build (MI355X, Triton 3.8.0): - add BF16 (4-wave, same 4096x4096x8192 shape as FP16); - add an 8-wave row for each data type at the matching shape; - refresh the FP16/BF8/MXFP4 4-wave numbers to the current build so the 4-wave vs 8-wave comparison is same-build (values match the -8wave READMEs); - note that BF16 measures ~5-6% above FP16 (clock/power effect; identical MFMA rate) and point to the pinned-build 1489 headline in a16w16/. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Lixun Zhang <lixun.zhang@amd.com>
…ne to v0 The 8-wave a4w4 B scale was byte-shuffled: at 8 warps (WARPS_N=4) an N-sliced [128,8] B-scale half gives each thread only 4 bytes, below the 64-bit ds_read_b64_tr_b8 transpose-read width, so it degraded to ds_read_u8 + v_perm (118 v_perm in the loop). v1_combineBsc loads the full [256,8] B scale as ONE combined buffer (8 bytes/thread -> transpose read, no v_perm) and splits it back into the left/right MFMA columns with a free register split + convert_layout. The combined [256,8] async fill needs a coalescing blocked layout ([4,1],[64,1],[1,8]: each warp = 64 N-lanes x 1 K-lane = one contiguous 256-byte K-column), because b_scales is N-contiguous in HBM ((K/32,N).T) and gfx950 direct-to-LDS cannot scatter (each warp's dword writes must be one contiguous LDS run). v0's [4,1],[32,2],[2,4] leaves a 128-byte gap between the two K groups a warp spans, which fails canCoalesceWriteIntoSharedMemory for [256,8]. Result (MI355X, rocprof cold-rotating): loop MFMA eff ~57% -> ~73%, TFLOPS +15-19% (K=32768: 4064 -> 4840), VGPR/spills 256/23 -> 256/12. Correct vs dequantized torch across K=1024..65536. - rename v1_sliceMN_BK256_nS2 -> v0_sliceMN_BK256_nS2 (the byte-shuffle baseline) - add v1_combineBsc_BK256_nS2 (default / recommended) - refresh a4w4-8wave + gemm READMEs with the v0-vs-v1 comparison and rationale Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Lixun Zhang <lixun.zhang@amd.com>
Explain why ds_read_b128 stalls ~21 cyc in the 8-wave scaled loop despite SQ_LDS_BANK_CONFLICT=0: a scaled MFMA hides a 4-cyc ld_scale that eats one of the two ds co-issue slots, and the SIMD-pair shared ds bus then serializes SIMD0/SIMD2 when they run in phase. 4-wave dodges it via a front-loaded ds that self-creates the skew; unscaled MFMA has two slots. Adds four hand-drawn timeline figures (scaled anatomy, 8-wave SP conflict, 4-wave auto-skew, non-scale two-slot) plus the scaled/non-scale ATT trace crops. Corroborated by the scaled->unscaled asm swap (ds stall 21.2->11.7). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Redraw the four schematic figures for clarity: each MFMA is now a staircase of 4-cyc boxes so the ld_scale/compute overlap is explicit (you see MFMA1's ld_scale sitting under MFMA0's compute); add a "1 MFMA = N boxes" brace and a "each box = 4 cycles" key; move the legend to a vertical block on the right (dropping unused swatches per figure); split ds markers above SIMD0 / below SIMD2; drop the cycle axis; shorten captions. Content unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- enlarge the legend and caption text; pull fig1's legend closer to the plot - add dashed vertical gridlines at box boundaries so same-cycle boxes of SIMD0 and SIMD2 line up (makes the SP-bus collision/skew legible) - fix fig3: SIMD0's ds markers were clipped at the top - drop the unused ld_scale swatch from fig4's legend Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Bump the §5 <img> display widths to 100% (schematic figures + the wide scaled ATT crop) and 66% for the narrower non-scale ATT crop, so the graphs read at the same scale as the surrounding text. Images unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Tighten each figure's lower y-limit to drop the empty band below the timeline and raise the caption toward the content, so the caption sits just under the graph. Nudge fig2's top limit back so SIMD0's Case-1 ds markers aren't clipped. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ads on v1 v2 replaces the earlier 32x32x64 experiment with the version that pairs the 32x32x64 MFMA with a width-matched, bank-conflict-free LDS/global-load layout (M-base reorder + [[1024,16]] padding). That drops SQ_LDS_BANK_CONFLICT 7.5M->1.2M and lifts MFMA occupancy ~81% -> ~98% (spill-free). The occupancy win is cycle-based, not wall-clock: v2 runs the loop in ~20% fewer cycles but the bigger MFMAs frequency-throttle ~21%, so TFLOPS is ~even with v1 (16x16x128). README covers the layout, the co-issue window, and the clock-throttle caveat. Also enable the fence_loads warp-pipeline flag on v1's mem stages, and register v2 as --version 2. NOTE: fence_loads needs a compiler prototype not in the pinned gfx950-tutorial-v0.3; these kernels require a Triton built with it. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…FP4 8-wave into v1 (4921/79.5%) and v2 (4750/98.0%)
…build (PR #10840) - a4w4-8wave v1 fence-on: ~73% -> ~80% loop MFMA (dir/v0/v1/v2 READMEs, gemm §1 + §5) - refresh §5 a16w16/a8w8 rows on the v1.0 build; FP16 8-wave now edges the 4-wave - §1 intro: cite gfx950-tutorial-v1.0 + §2.1, drop stale env vars / Triton 3.8.0; NOTE ~5-6% -> ~6% - naming: llir+amdgcnas -> llir+force-agpr+amdgcnas
… table - Versions table: add v2 row; describe v2 as the ~98%-MFMA, spill-free, clock-throttled cycle-efficiency variant (wall-clock even with v1) - §3.1: v2-vs-v1 perf table (K=8192/32768) - §3.2: fence_loads A/B (v1 +6.4pp, v2 +7.2pp loop MFMA at K=32768) - Running/Files: add --version 2 commands and the v2 directory
The a16w16-8wave / a8w8-8wave / a4w4-8wave Python files were added without running black/ruff, so the python-format-check CI failed on the branch. Apply black (line-length 100) and ruff import sorting (I001). Formatting only — no logic changes.
8fa9d62 to
217b94e
Compare
fa0d581 to
690230e
Compare
…ve/ (8-wave)
Move the six gemm kernel dirs into two scheduling-model groups:
a16w16/a8w8/a4w4 -> intra_wave/{a16w16,a8w8,a4w4}
a16w16-8wave/a8w8-8wave/... -> inter_wave/{a16w16,a8w8,a4w4}
Update all path-dependent references: run_perf_table.py / run_counter_collection.py
work-dirs, the 4-wave bench.py plugin sys.path depth, every relative markdown link
(root/gemm-README/cross-group), the gemm README structure (tree, Versions table,
§1/§5), top README, and docs. Verified: all links resolve, black+ruff pass, and
4-wave (full stack) + 8-wave kernels still pass correctness after the move.
- common.py (inter_wave a16w16/a8w8/a4w4): the store epilogue mask compared tile-local offs_cm/offs_cn (0..BLOCK-1) against the full M/N, so it was always true and never masked the ragged last tile. Add the pid*BLOCK tile offset so non-divisible M/N are masked correctly. Only a16w16 v0 calls store_result today; the shipped 4096x4096 benchmarks are unaffected. - scripts/README.md: update the ATT profiling example for the intra_wave/ restructure (cd kernels/gemm/intra_wave/a16w16 + ../../../../scripts). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds an 8-wave warp-pipeline variant of each tutorial GEMM —
a16w16-8wave(FP16/BF16),a8w8-8wave(BF8), anda4w4-8wave(MXFP4) — as a second route to high MFMA utilization alongside the 4-wavellir+force-agpr+amdgcnaskernels onmain.Instead of the LLIR scheduler + amdgcnas, these launch 8 warps/CTA (2 waves/SIMD) and schedule the hot loop at the wave level with
warp_pipeline_stage: the two resident waves per SIMD are kept out of phase (a "ping-pong") so one issues MFMAs while the other issues loads, then they swap. They run with no AGPRs (amdgpu-agpr-alloc=0,0viallvm_fn_attrs) — the f32 accumulators live in VGPRs — so no environment variables are needed.What's in here
Kernels
a16w16-8wave/(FP16/BF16) —v0_BK32_nS3,v1_sliceMN_BK64_nS2a8w8-8wave/(BF8) —v1_sliceMN_BK128_nS2a4w4-8wave/(MXFP4) —v0_sliceMN_BK256_nS2(byte-shuffle B scale),v1_combineBsc_BK256_nS2(combined transpose-read B scale, recommended),v2_mfma32x32x64_BK256_nS2(32×32×64 MFMA + conflict-free LDS layout)Each directory ships
bench.py(correctness + rocprof rotating-tensor timing),collect_perf.py(TFLOPS + ATT MFMA efficiency + VGPR/spill), and per-version READMEs.Docs
docs/warp_pipelining.md— the warp-pipelining theory reference.kernels/gemm/README.md §5— an "8-Wave Warp-Pipeline Variants" section: the design, a kernel comparison, a cross-K performance table, and where the 8-wave lands vs the 4-wave. §1's summary table gains BF16 and 8-wave rows.fence_loadsdependencyThe
a4w4-8wavev1 and v2 kernels usewarp_pipeline_stage(..., fence_loads=True)— a fullsched.barrieremitted after the stage's LDS reads, which stops the backend from hoisting the next tile's global→LDS prefetch ahead of the reads a following MFMA depends on. This needs a Triton withfence_loadssupport (triton-lang/triton#10840), which is not in the pinnedgfx950-tutorial-v1.0tag — those two kernels require a Triton built with it. The other 8-wave kernels do not use the fence and build againstv1.0directly.Performance highlights (MI355)
v1edges the 4-wavev9(~1442 vs 1421 TFLOPS @ K=8192, ~99.8% loop MFMA).fence_loadsintra-stage fence lifts v1's loop MFMA from ~73% → ~80%; v2 (32×32×64 + conflict-free layout) reaches ~98% occupancy but is a wall-clock wash with v1 (clock throttle on the wider MFMAs).Full numbers are in each
-8wave/README.mdandgemm/README §1/§5.Notes
llir+force-agpr+amdgcnastoolchain is built around the 4-wave register/schedule model and fails register allocation at 8 waves, so it is not used here.python-format-checkpasses.🤖 Generated with Claude Code